#' ----------------------------------------------------------------------------
#' Comparing Estimations of Bayesian Aldrich-Mckelvey Models
#' 
#' Description: This script extracts the main parameters of interests from the 
#' BAM models and compare them.

# Install and load required packages -------------------------------------------
if (!require("pacman")) install.packages("pacman")
p_load(tidyverse, rstan, hbamr)

# Load data --------------------------------------------------------------------

# BAM Estimates
bam_fit <- readRDS("outputs/bam.rds")
hbam_fit <- readRDS("outputs/hbam_base.rds")
hbam_multi_fit <- readRDS("outputs/hbam_multi.rds")

# CHES Data
df_expert <- readRDS("outputs/df_global_expert.rds")
df_expert_party <- readRDS("outputs/df_global_expert_party.rds")
df_bam <- readRDS("outputs/df_bam.rds")

# Functions --------------------------------------------------------------------
# Function to extract parameters of interest and estimate quantiles.
ex_fun <- function(data, quan) {
  df <- as.data.frame(rstan::extract(data, quan)) %>%
    sapply(., quantile, probs = c(0.025, 0.5, 0.975)) %>%
    t() %>%
    as.data.frame()
  return(df)
}

# Comparing Party Positions ----------------------------------------------------

# Extract Party Positions
bam_pp <- ex_fun(bam_fit, "Z") %>% mutate(pnum = seq(1:nrow(.)))
hbam_pp <- ex_fun(hbam_fit, "theta")
hbam_m_pp <- ex_fun(hbam_multi_fit, "theta")

# Combine HBAM estimates in a single dataframe
df_hbam_pp <- cbind(hbam_pp, hbam_m_pp)[, c(2, 5)] # Only 0.5 percentile
names(df_hbam_pp) <- c("hbam", "hbam_m")

# Add CHES data to HBAM dataset
df_hbam_pp <- df_hbam_pp %>% 
  mutate(party_id = gsub(".*_", "", names(df_expert[8:ncol(df_expert)])))

# Extract Party IDs and Pnum and merge with CHES data
df_bam_pnum <- merge(
  df_bam %>% 
    group_by(pnum) %>% 
    summarise(
      party_id = first(party),
      party_id = gsub("lrecon_", "", party_id)
    ),
  df_expert_party %>% 
    group_by(region_id, country, party_id, party) %>%
    summarise(lrecon = mean(lrecon, na.rm = TRUE)),
  by = "party_id"
)

# Merge BAM Party ID and Pnum with estimations/with HBAM
## BAM
df_bam_pp <- merge(
  df_bam_pnum, 
  bam_pp %>% 
    select("50%", "pnum") %>% 
    rename("bam" = "50%"), # Only 0.5 percentile
  by = "pnum"
)
## HBAM
df_all <- merge(df_bam_pp, df_hbam_pp, by = "party_id")

# TABLE A1: Correlation of estimates
cor(df_all[c(6:ncol(df_all))])

# Delete residual objects
rm(bam_pp, hbam_pp, hbam_m_pp, df_bam_pnum, df_bam_pp, df_hbam_pp, df_all)

# ALPHA/BETAS ------------------------------------------------------------------

# HBAM
## Extract Alphas and Betas
hbam_ab <- cbind(
  ex_fun(hbam_fit, "alpha"), 
  ex_fun(hbam_multi_fit, "alpha"),
  ex_fun(hbam_fit, "beta"),
  ex_fun(hbam_multi_fit, "beta")
)[, c(2, 5, 8, 11)] %>% # Only 0.5 percentile 
  rename(
    "h_alpha" = 1, 
    "hm_alpha" = 2, 
    "h_beta" = 3, 
    "hm_beta" = 4
  )

## Merge with Expert-level data
df_hbam_ab <- cbind(df_expert[, 1:7], hbam_ab)
  
# BAM
## Extract and add observation numbers
bam_ab <- cbind(
  ex_fun(bam_fit, "alpha"), 
  ex_fun(bam_fit, "beta")
)[, c(2, 5)] %>% 
  mutate(obs = row_number()) %>% 
  rename(
    "b_alpha" = 1, 
    "b_beta" = 2
  )

## Merge data
df_bam_ab <- merge(
  df_bam %>% 
    group_by(expert_id) %>% 
    summarise(obs = first(obs)), 
  bam_ab, 
  by = "obs"
)

# Merge all
df_ab <- merge(df_bam_ab, df_hbam_ab, by = "expert_id")

# Re-order and select variables of interest
df_ab <- df_ab %>%
  select(
    b_alpha,
    h_alpha,
    hm_alpha,
    b_beta,
    h_beta,
    hm_beta
  )

# Correlations
## Table A1: Alpha
cor(df_ab[, c(1:3)])
## Table A2: Beta
cor(df_ab[, c(4:6)])

# Save session information for reproducibility
writeLines(capture.output(sessionInfo()), "outputs/session_info/03_correlations.txt")